# -*- coding: utf-8 -*-
"""
Created on Fri Feb  3 11:46:18 2023

@author: 潘宇浩
"""


def read_record():
    fr=open("record.txt","r")
    lines=fr.readlines()
    x_cal=[]
    y_cal=[]
    x_cal.append([float(lines[-1].split()[0]),float(lines[-1].split()[1])])
    y_cal.append(0)
    best_Ener=float(lines[-1].split()[4])
    for i in range(1,len(lines)-2):
        x_cal.append([float(lines[i].split()[0]),float(lines[i].split()[1])])
        y_cal.append(round(float(lines[i].split()[4])-best_Ener,5))      
    return x_cal,y_cal

def read_input():
    input_f=open('input','rt')
    input_lines=input_f.readlines()
    x_lb=[float(input_lines[3]),float(input_lines[4])]
    x_ub=[float(input_lines[5]),float(input_lines[6])]
    return x_lb,x_ub

def draw_ea(x_cal,y_cal,x_lb,x_ub,draw_type='N'):
    from scipy.interpolate import griddata
    import numpy as np
    import matplotlib.pyplot as plt
    for i in range(len(x_cal)):
        if x_cal[i][0] == 0:
            for j in range(11):
                x_cal.append([0,x_lb[1]+(x_ub[1]-x_lb[1])/10*j])
                y_cal.append(y_cal[i])
    if draw_type=='i':
        for j in range(len(x_cal)):
            x_cal.append([180-x_cal[j][0],x_cal[j][1]])
            y_cal.append(y_cal[j])
    theta_grid, phi_grid = np.mgrid[x_lb[0]:x_ub[0]:1001j, x_lb[1]:x_ub[1]:1001j]
    y_grid=griddata(np.array(x_cal), np.array(y_cal), (theta_grid, phi_grid),method='cubic')
    plt.switch_backend('agg')
    plt.contourf(theta_grid, phi_grid, y_grid,100,)
    plt.colorbar()
    plt.savefig('EMA_mesh.png')

def write_ea(x_cal,y_cal):
    fr=open("out.txt","wt")
    fr.write('theta(°)\tphi(°)\tEnergy(eV)')
    for i in range(len(y_cal)):
        fr.write("\n%.3f\t%.3f\t%.8f"%(x_cal[i][0],x_cal[i][1],y_cal[i]))
    fr.close()
    return

if __name__=='__main__' :
    import sys
    x_cal,y_cal=read_record()
    write_ea(x_cal, y_cal)
    x_lb,x_ub=read_input()
    if len(sys.argv) <=1:
        draw_ea(x_cal,y_cal,x_lb,x_ub)
    else:
        if sys.argv[1] == "i":
            draw_ea(x_cal,y_cal,x_lb,x_ub,draw_type='i')